'''
Created on 28 Jun 2013

@author: Matt
'''

from Crypto.Cipher import AES
import binascii
from Crypto.Hash import SHA256
import sys
import os

#Takes a hex character and converts it into a 4 bit array 
#representing the bits that make up the hex character
def hexToBits(hexChar):
    if (hexChar == '0'):
        return [0,0,0,0]
    elif (hexChar == '1'):
        return [0,0,0,1]
    elif (hexChar == '2'):
        return [0,0,1,0]
    elif (hexChar == '3'):
        return [0,0,1,1]
    elif (hexChar == '4'):
        return [0,1,0,0]
    elif (hexChar == '5'):
        return [0,1,0,1]
    elif (hexChar == '6'):
        return [0,1,1,0]
    elif (hexChar == '7'):
        return [0,1,1,1]
    elif (hexChar == '8'):
        return [1,0,0,0]
    elif (hexChar == '9'):
        return [1,0,0,1]
    elif (hexChar == 'A'):
        return [1,0,1,0]
    elif (hexChar == 'B'):
        return [1,0,1,1]
    elif (hexChar == 'C'):
        return [1,1,0,0]
    elif (hexChar == 'D'):
        return [1,1,0,1]
    elif (hexChar == 'E'):
        return [1,1,1,0]
    elif (hexChar == 'F'):
        return [1,1,1,1]

#Iterates over an entire hex string returning a list of bits
#that represent the binary versino of the hex string
def hexStrToBits(hexStr):
    hexbytes = []
    i = 0
    while i < len(hexStr)-1:
        #print("Getting hexbits for: " + hexStr[i] + hexStr[i+1])
        chars1 = hexToBits(hexStr[i].upper());
        for c in chars1:
            hexbytes.append(c)
        i = i + 1
        chars2 = hexToBits(hexStr[i].upper());
        #print("Got: " + str(chars1) + str(chars2))
        for c in chars2:
            hexbytes.append(c)
        i = i + 1
        
    return hexbytes

def hexor(str1, str2):
    bits1 = hexStrToBits(str1)
    bits2 = hexStrToBits(str2)
    bitsXor = xor(bits1, bits2)
    strXor = bitsToHex(bitsXor)
    return strXor

#xors two lists of bits, returning the xor 
#of the min length of the two lists
def xor(str1, str2):
    xorbytes = []
    length = min(len(str1), len(str2))
    for i in range(length):
        if str1[i] == str2[i]:
            xorbytes.append(0)
        else:
            xorbytes.append(1)
            
    return xorbytes

#Takes an ascii string and returns the list of bits 
#representing that ascii string
def strToBits(str):
    result = []
    for c in str:
        bits = bin(ord(c))[2:]
        bits = '00000000'[len(bits):] + bits
        result.extend([int(b) for b in bits])
    return result 

def intToBits(n):#
    return [int(digit) for digit in bin(n)[2:]] 
    
#takes a bytes object and returns the list of
#bits representing those bytes
def bytesToBits(str):
    result = []
    for c in str:
        bits = bin(c)[2:]
        bits = '00000000'[len(bits):] + bits
        result.extend([int(b) for b in bits])
    return result        


def AES_CBC_Decrypt(hexEncodedIV, hexEncodedCipher, hexEncodedKey):
    myAES = AES.new(bytes.fromhex(hexEncodedKey))
    myAES.mode = AES.MODE_ECB
    c = ["" for x in range(len(hexEncodedCipher) // 32)]
    for x in range(0, len(hexEncodedCipher) // 32):
        start = (x * 32)
        end = ((x + 1) * 32)
        c[x] = hexEncodedCipher[start:end]
        
    m = ["" for x in range(0, len(c))]
    plainStr = ""
    for x in range(0, len(c)):
        d = bytesToBits(myAES.decrypt(bytes.fromhex(c[x])))
        if x == 0:
            m[x] = xor(d, hexStrToBits(hexEncodedIV))
        else:
            m[x] = xor(d, hexStrToBits(c[x-1]))
        binStr = "0b" + (("").join(map(str, m[x])))
        n = int(binStr, 2)
        plain = binascii.unhexlify('%x' % n)
        plainStr += str(plain)[2:-1]
    intValOfLastByte = int(plainStr[-2:], 16)
    plainStr = plainStr[0:(intValOfLastByte*-4)]
    print(plainStr)
    
def bitsToHex(listOfBits):
    binStr = (("").join(map(str, listOfBits)))
    n = hex(int(binStr, 2))
    return n[2:]
    
def AES_CTR_Decrypt(hexEncodedIV, hexEncodedCipher, hexEncodedKey):
    myAES = AES.new(bytes.fromhex(hexEncodedKey))
    myAES.mode = AES.MODE_ECB
    c = ["" for x in range((len(hexEncodedCipher) // 32) + 1)]
    for x in range(0, (len(hexEncodedCipher) // 32) + 1):
        start = (x * 32)
        end = ((x + 1) * 32)
        c[x] = hexEncodedCipher[start:end]
        
    m = ["" for x in range(0, len(c))]
    plainStr = ""
    for x in range(0, len(c)):
        d = bytesToBits(myAES.encrypt(bytes.fromhex("%x" % (int(hexEncodedIV, 16) + x))))
        m[x] = xor(d, hexStrToBits(c[x]))
        binStr = "0b" + (("").join(map(str, m[x])))
        n = int(binStr, 2)
        plain = binascii.unhexlify('%x' % n)
        plainStr += str(plain)[2:-1]
    print(plainStr)
        
def hash_file_in_blocks(blockSize, fileName):
    file = open(fileName, "rb")
    sizeOfFile = os.stat(fileName).st_size
    numBlocks = sizeOfFile // blockSize
    #print(sizeOfFile)
    currentBlock = numBlocks
    lastHash = b""
    #print(numBlocks)
    try:
        while currentBlock > -1:
            sha = SHA256.new()
            file.seek(currentBlock * blockSize)
            bytes_read = b""
            bytes_read += file.read(blockSize)
            #print(str(currentBlock) + ": " + str(len(bytes_read)))
            bytes_read += lastHash
            #print(str(currentBlock) + ": " + str(len(bytes_read)))
            sha.update(bytes_read)
            lastHash = sha.digest()
            #print(binascii.hexlify(lastHash))
            currentBlock -= 1
        #while bytes_read:
        #    sha.update(bytes_read)
        #    print(sha.digest())
        #    bytes_read = file.read(blockSize)
    finally:
        file.close()
    
    return binascii.hexlify(lastHash)


BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS) 
unpad = lambda s : s[0:-4*int(s[-1])]
    
